from typing import List, Dict, Any
from base import Task,Job,create_job

class ParamterGenerator:
    '''
    根据不同任务，生成对应的模型参数
    '''

    def __init__(self, task: Task):
        """
        初始化：直接接收Task对象，从中提取工件信息
        """
        # 从任务对象中提取所需信息，不直接依赖外部零散参数
        self.job_types: List[Job] = task.job_types
        self.job_counts: List[int] = task.job_counts
        self.K: List[int] = task.K
        self.R: List[int] = task.R

        # 计算所有工件全局工序总数
        self.total_ops_num = sum(
            # 每个工件工序数 * 工件数 = 工序总数
            job.ops_nums * count
            for job, count in zip(self.job_types, self.job_counts)
        )

        self.J = list(range(1, self.total_ops_num + 1))  # 全局工序编号[1..总工序数]
        # 计算各类型的工件的“起始工序偏移量”
        self.job_type_offsets = self._precompute_type_offsets()

    def _precompute_type_offsets(self) -> List[int]:
        """
        计算每种工件类型的“起始工序偏移量”（按索引对应）：
        返回列表，索引与job_types一致，值为该类型第一道工序的【全局起始编号,从0开始】
        示例：A类型2个（2道工序）、B类型3个（3道工序）→ 偏移量[0, 4]
        """
        type_offsets = []
        cumulative = 0  # 累计前面类型的总工序数
        for job, count in zip(self.job_types, self.job_counts):
            type_offsets.append(cumulative)
            # 累加当前类型的总工序数（数量×单工件工序数）
            cumulative += job.ops_nums * count
        return type_offsets

    def get_type_info(self, op_id: int) -> (Job, int, int):
        """
        获取全局工序 op_id 所属的工件类型、在单个工件内的步骤、以及该工件在类型内的索引
        返回：(工件类型, 单个工件工序编号, 单个工件索引)
        """
        if op_id < 1 or op_id > self.total_ops_num:
            raise ValueError(f"工序 {op_id} 超出全局工序范围（1~{self.total_ops_num}）")

        cumulative = 0  # 累计前面所有类型的总工序数（用于定位当前类型的工序范围）
        # 遍历“类型-数量”对应关系（索引对齐）
        for type_idx in range(len(self.job_types)):
            job_type = self.job_types[type_idx]  # 当前工件类型
            type_count = self.job_counts[type_idx]  # 当前类型工件的实际工件数量
            single_op_num = job_type.ops_nums  # 当前类型单个工件的工序数
            type_total_ops = single_op_num * type_count  # 当前类型的总工序数

            # 判断当前工序op_id是否在当前类型的工序范围内
            if cumulative <= op_id - 1 < cumulative + type_total_ops:
                # 1. 计算当前全局工序在 “当前类型工件内的相对编号”（从1开始：1 2 3 4...）
                type_internal_id = op_id - cumulative
                # 2. 计算该工序所属的“当前类型内的工件索引”（从0开始）
                # 逻辑：(相对编号-1) // 单工件工序数 → 整除得到工件索引
                workpiece_idx = (type_internal_id - 1) // single_op_num
                # 3. 计算该工序在“单个工件内的步骤”（从1开始）
                # 逻辑：(相对编号-1) % 单工件工序数 + 1 → 取余得到步骤
                step = (type_internal_id - 1) % single_op_num + 1
                return job_type, step, workpiece_idx

            # 累加当前类型的总工序数，进入下一个类型的判断
            cumulative += type_total_ops

        # 若遍历完所有类型仍未找到，说明工序编号异常
        raise ValueError(f"工序 {op_id} 未找到对应工件类型")

    def generate_Ns(self) -> List[List[int]]:
        """
        生成Ns：同一类型工件的相同工艺步骤工序集合
        示例（A类型2个+ B类型3个）：[[1,3], [2,4], [5,8,11], [6,9,12], [7,10,13]]
        """
        Ns: List[List[int]] = []
        # 按索引遍历：同一索引对应“一种类型+其数量”
        for type_idx in range(len(self.job_types)):
            job_type = self.job_types[type_idx]  # 当前类型
            count = self.job_counts[type_idx]  # 当前类型的数量
            start_offset = self.job_type_offsets[type_idx]  # 当前类型的起始偏移量

            # 遍历该类型的每个工序步骤（1~wt.ops_nums）
            for step in range(1, job_type.ops_nums + 1):
                step_ops = []
                # 遍历该类型的每个工件（0~count-1）
                for workpiece_idx in range(count):
                    # 计算全局工序编号：偏移量 + 工件索引×单工件工序数 + 步骤号
                    op_number = start_offset + workpiece_idx * job_type.ops_nums + step
                    step_ops.append(op_number)
                Ns.append(step_ops)
        return Ns

    def generate_Nj(self) -> List[List[int]]:
        """
        生成Nj：每个工件的所有工序集合
        例如：A类型2个（各2道工序）+ B类型1个（3道工序）→
        [[1,2], [3,4], [5,6,7]]
        """
        Nj: List[List[int]] = []

        for type_idx in range(len(self.job_types)):
            job_type = self.job_types[type_idx]  # 当前工件类型
            count = self.job_counts[type_idx]  # 该类型的工件数量
            start_offset = self.job_type_offsets[type_idx]  # 类型起始偏移量
            ops_per_job = job_type.ops_nums  # 每个工件的工序数

            # 遍历该类型的每个工件
            for workpiece_idx in range(count):
                # 计算当前工件的首道工序编号
                job_start = start_offset + workpiece_idx * ops_per_job + 1
                # 计算当前工件的末道工序编号
                job_end = job_start + ops_per_job - 1
                # 生成当前工件的所有工序编号
                job_ops = list(range(job_start, job_end + 1))
                Nj.append(job_ops)

        return Nj

    def generate_Nf(self) -> List[int]:
        """
        生成Nf：每个工件的首工序（索引对齐）
        示例（A类型2个+ B类型3个）：[1,3,5,8,11]
        """
        Nf: List[int] = []
        for type_idx in range(len(self.job_types)):
            wt = self.job_types[type_idx]
            count = self.job_counts[type_idx]
            start_offset = self.job_type_offsets[type_idx]

            for workpiece_idx in range(count):
                # 首工序 = 偏移量 + 工件索引×单工件工序数 + 1
                first_op = start_offset + workpiece_idx * wt.ops_nums + 1
                Nf.append(first_op)
        return Nf

    def generate_Nl(self) -> List[int]:
        """
        生成Nl：每个工件的尾工序（索引对齐）
        """
        Nl: List[int] = []
        for type_idx in range(len(self.job_types)):
            wt = self.job_types[type_idx]
            count = self.job_counts[type_idx]
            start_offset = self.job_type_offsets[type_idx]
            single_ops_num = wt.ops_nums  # 单工件工序数

            for workpiece_idx in range(count):
                # 尾工序 = 偏移量 + (工件索引+1)×单工件工序数
                last_op = start_offset + (workpiece_idx + 1) * single_ops_num
                Nl.append(last_op)

        return Nl


    def generate_PM_PP(self) -> (Dict[tuple, int], Dict[tuple, int]):
        """
        生成加工可行性（PM）和加工时间（PP）
        PM[(j, k)]: 1=工序j可在机器k上加工，0=不可加工
        PP[(j, k)]: 工序j在机器k上的加工时间（不可加工时设为1e10）
        适配多类型工件：依赖get_type_info获取工序所属类型及步骤，补充异常处理确保鲁棒性
        """
        PM = {}
        PP = {}
        # 定义“不可加工”的默认加工时间（避免魔法数字）
        INFEASIBLE_TIME = 1e10

        for j in self.J:
            try:
                # 1. 获取工序j所属的类型、类型内步骤
                wt, step, _ = self.get_type_info(j)

                # 2. 校验当前类型的工艺规则是否完整
                if step not in wt.feasible_rules:
                    raise KeyError(f"工件类型[{wt.name}]的工序步骤[{step}]未配置可行机器规则")
                if step not in wt.process_time_rules:
                    raise KeyError(f"工件类型[{wt.name}]的工序步骤[{step}]未配置加工时间规则")
                # 校验加工时间是否为非空列表（避免索引越界）
                if not wt.process_time_rules[step] or len(wt.process_time_rules[step]) == 0:
                    raise ValueError(f"工件类型[{wt.name}]的工序步骤[{step}]加工时间配置为空")

                # 3. 遍历所有机器，判断加工可行性并赋值
                for k in self.K:
                    if k in wt.feasible_rules[step]:
                        # 可加工：PM=1，PP=配置的加工时间
                        PM[(j, k)] = 1
                        PP[(j, k)] = wt.process_time_rules[step][0]
                    else:
                        # 不可加工：PM=0，PP=默认不可行时间
                        PM[(j, k)] = 0
                        PP[(j, k)] = INFEASIBLE_TIME

            except Exception as e:
                # 捕获异常并明确报错信息
                raise RuntimeError(f"生成工序[{j}]的PM/PP时出错：{str(e)}") from e

        return PM, PP



    def generate_TM(self) -> Dict[tuple, int]:
        """
        生成运输可行性（TM）：TM[(j, r)]=1 表示工序j可由机械手r运输，0=不可运输
        """
        TM = {}
        for j in self.J:
            try:
                # 1. 获取工序j所属的工件类型、单个工件内的类型内步骤
                job, step, _ = self.get_type_info(j)

                # 2. 校验当前类型的运输规则是否完整（避免KeyError）
                if step not in job.transport_rules:
                    raise KeyError(f"工件类型[{job.name}]的工序步骤[{step}]未配置运输机械手规则")
                # 校验运输规则是否为非空列表（避免无效配置）
                if not job.transport_rules[step] or len(job.transport_rules[step]) == 0:
                    raise ValueError(f"工件类型[{job.name}]的工序步骤[{step}]运输机械手配置为空")

                # 3. 遍历所有机械手，判断运输可行性并赋值
                for r in self.R:
                    TM[(j, r)] = 1 if r in job.transport_rules[step] else 0

            except Exception as e:
                # 封装异常信息，便于定位多类型场景下的配置问题
                raise RuntimeError(f"生成工序[{j}]的运输可行性（TM）时出错：{str(e)}") from e

        return TM

    def generate_TT(self) -> Dict[tuple, int]:
        """
        生成运输时间（TT）：TT[(k_prev, k_curr)] 表示从机器k_prev到k_curr的运输时间
        """
        TT = {}

        # 遍历所有机器对，生成运输时间
        for k_prev in self.K:
            for k_curr in self.K:
                if k_prev == k_curr:
                    TT[(k_prev, k_curr)] = 0
                else:
                    TT[(k_prev, k_curr)] = 2

        return TT

    def generate_job_mapping(self) -> Dict[int, int]:
        """
        生成“工序号→工件ID”的映射字典
        返回：key=全局工序号，value=工件ID
        示例：Nj=[[1,2],[3,4]] → job_mapping={1:1, 2:1, 3:2, 4:2}
        """
        job_mapping = {}
        # Nj 中每个子列表对应一个工件，索引+1即为工件ID
        for job_id, job_ops in enumerate(self.generate_Nj(), start=1):
            # 遍历当前工件的所有工序，赋值映射
            for op_id in job_ops:
                job_mapping[op_id] = job_id
        return job_mapping

    def generate_all_params(self) -> Dict[str, Any]:
        """生成所有参数"""
        # 加工可行性和加工时间
        PM, PP = self.generate_PM_PP()
        # 运输可行性
        TM = self.generate_TM()
        # 运输时间（简化为固定值）
        TT = self.generate_TT()

        return {
            'J': self.J,
            'K': self.K,
            'R': self.R,
            'Ns': self.generate_Ns(),
            'Nj': self.generate_Nj(),
            'Nf': self.generate_Nf(),
            'Nl': self.generate_Nl(),
            'job_mapping': self.generate_job_mapping(),
            'M': 1e6,
            'alpha': 15,
            'beta': 30,
            'PM': PM,
            'PP': PP,
            'TM': TM,
            'TT': TT,
        }

    def pretty_print_dict(self,Paramters, indent=0):
        """递归格式化字典，支持缩进和分层"""
        for key, value in Paramters.items():
            # 打印当前层级的键
            print("  " * indent + f"{key}: ", end="")
            if isinstance(value, dict):
                # 如果值是字典，递归缩进
                print("{")
                self.pretty_print_dict(value, indent + 1)
                print("  " * indent + "},")
            elif isinstance(value, list):
                # 如果值是列表，简单格式化
                print(f"{value},")
            else:
                # 其他类型直接打印
                print(f"{value},")



# ------------------ 示例验证 ------------------
if __name__ == "__main__":
    # 任务信息
    # 创建工件
    job_creator = Job.create_job()
    job_d = job_creator.generator_job_d()
    job_c = job_creator.generator_job_c()

    # 创建任务实例
    task = Task(
        name="任务-5c",
        job_types=[job_c, job_d],
        job_counts=[25,50]
    )

    task.print_task_info()

    paramter_generator = ParamterGenerator(task)

    # 生成参数
    params = paramter_generator.generate_all_params()
    paramter_generator.pretty_print_dict(params)


